import pymysql
import datetime
import json
import subprocess

"""
清空表空间后cp复制frm和ibd文件,再导入表空间,来恢复数据;
用于表批量备份后的恢复;
修改mysql连接参数和目录路径参数后执行;
"""

# python执行shell
def run_cmd(shell):
    try:
        cp = subprocess.run(shell, shell=True, capture_output=True, encoding='utf-8')
    except Exception as ex:  # 如果 try 段代码有出现异常,则执行
        print(ex)
    else:
        time_stamp = datetime.datetime.now().strftime('%Y.%m.%d-%H:%M:%S')
        print(time_stamp + "    " + shell)
        return cp.stdout, cp.stderr


# 读取txt文件,整理后转换为json文件
def txt_to_json(txt_path):
    dic_database_table = {}  # 定义一个空的字典

    f = open(txt_path, 'r')
    lines = f.readlines()
    for line in lines:
        new_lines = line.strip()  # 去除空行、空格
        data_name = new_lines.split(".")[0]  # 获取这一行中的数据库名称
        table_name = new_lines.split(".")[1]  # 获取这一行中的表名称

        # 1-如果数据库名称存在字典的key中，则将通过 字典[数据库名称].append 的方式添加进空数组中；
        if data_name in dic_database_table:
            dic_database_table[data_name].append(table_name)
        # 2-如果数据库名称不存在字典的key中，则报错；
        else:
            dic_database_table[data_name] = []  # 如果不存在这个key,则说明是第一次,就建立一个空的数组
    f.close()
    # print(dic_database_table)

    # 将python字典写入json文件
    # filename = "target_backupTable.json"  # 处理后要有生成的json文件
    # with open(filename, 'w') as f_obj:
    #     json.dump(dic_database_table, f_obj, indent=4, sort_keys=False)

    return dic_database_table


# 连接mysql执行sql的方法(需要手动修改 host/user/port/password 参数)
def mysql_cmd(mysql_name, sql_cmd):
    # 打开数据库连接
    db_connect = pymysql.connect(host="192.168.2.15", user="root", port=3306, password="heiBAI@xxzx#2020",
                                 db=mysql_name)
    cursor = db_connect.cursor()  # 使用cursor()方法获取操作游标
    # 执行sql语句
    try:
        cursor.execute(sql_cmd)
        db_connect.commit()  # 提交
    except Exception as ex:  # 如果 try 段代码有出现异常,则执行
        db_connect.rollback()  # 回滚
        print(ex)
        return False
    return True


def data_uiaer(data_name):
    list_discard = []  # 存储清空表错误时的命令
    list_cmd_cp = []  # 存储cp复制命令
    list_import = []  # 存储导入表的命令
    result_cmd_xpt = True

    # 0-关闭外键约束检查
    sql_close_key_checks = "SET Foreign_key_checks=0;"
    result_close_key_checks = mysql_cmd(data_name, sql_close_key_checks)
    time_stamp = datetime.datetime.now().strftime('%Y.%m.%d-%H:%M:%S')
    if result_close_key_checks:
        print(time_stamp + "    " + data_name + "---关闭外键约束检查---成功")
    else:
        print(time_stamp + "    " + data_name + "---关闭外键约束检查---失败")
        exit()

    for table_name in dic_database_table[data_name]:  # 遍历表
        # 定义要执行的命令
        sql_cmd_discard = "alter table " + data_name + "." + table_name + " discard tablespace;"
        cmd_cp_frm = "/usr/bin/cp -ar " + data_path + data_name + "/" + table_name + ".frm " + mysql_datadir + "/" + data_name + "/"
        cmd_cp_ibd = "/usr/bin/cp -ar " + data_path + data_name + "/" + table_name + ".ibd " + mysql_datadir + "/" + data_name + "/"
        cmd_chown = "chown -R " + docker_userGroup + " " + mysql_datadir
        sql_cmd_import = "alter table " + data_name + "." + table_name + " import tablespace;"

        # 1-清空表空间
        result_discard = mysql_cmd(data_name, sql_cmd_discard)
        time_stamp = datetime.datetime.now().strftime('%Y.%m.%d-%H:%M:%S')
        if result_discard:
            print(time_stamp + "    清空表空间成功---" + table_name)
        else:
            print(time_stamp + "    \033[31m清空表空间失败---%s\033[0m" % table_name)

            # 定义需要执行的命令(实际不执行)将需要执行的命令写入txt文本
            list_discard.append(sql_cmd_discard)
            list_cmd_cp.append(cmd_cp_frm)
            list_cmd_cp.append(cmd_cp_ibd)
            list_import.append(sql_cmd_import)

        # 2-复制表数据
        run_cmd(cmd_cp_frm)
        run_cmd(cmd_cp_ibd)
        # 3-修改数据库目录用户权限
        run_cmd(cmd_chown)
        # 4-导入表空间
        result_import = mysql_cmd(data_name, sql_cmd_import)
        time_stamp = datetime.datetime.now().strftime('%Y.%m.%d-%H:%M:%S')
        if result_import:
            print(time_stamp + "    导入表空间成功---" + table_name)
        else:
            print(time_stamp + "    \033[31m导入表空间失败---%s\033[0m" % table_name)

        if len(list_discard) != 0:
            result_cmd_xpt = False

    # 5-关闭外键约束检查
    sql_open_key_checks = "SET Foreign_key_checks=1;"
    result_open_key_checks = mysql_cmd(data_name, sql_open_key_checks)
    time_stamp = datetime.datetime.now().strftime('%Y.%m.%d-%H:%M:%S')
    if result_open_key_checks:
        print(time_stamp + "    " + data_name + "---恢复外键约束检查---成功")
    else:
        print(time_stamp + "    " + data_name + "---恢复外键约束检查---失败")
        exit()

    return list_discard, list_cmd_cp, list_import, result_cmd_xpt


if __name__ == "__main__":
    # 定义的参数---需要手动 修改
    mysql_datadir = "/usr/local/mysql/data_kongbiao"  # mysql的数据目录(不需要携带/)
    txt_path = "target_backupTable.txt"  # 定义txt文件路径
    data_path = "./data/"  # 原数据目录的路径(携带/)
    docker_userGroup = "polkitd:mysql"      # 指定的用户和用户组权限

    # 处理txt文件,将txt转换成json格式;
    dic_database_table = txt_to_json(txt_path)  # 将txt处理转换为字典格式
    # print(dic_database_table)

    judge_result_list = True  # 用于判断最后是否输出手动恢复的提示
    list_discard = []  # 存储清空表错误时的命令
    list_cmd_cp = []  # 存储cp复制命令
    list_import = []  # 存储导入表的命令

    echo_cmd = """echo "" > error_manualOperation.txt"""  # shell清空txt文件
    run_cmd(echo_cmd)
    file_error = open('error_manualOperation.txt', 'a')

    for data_name_1 in dic_database_table.keys():  # 遍历数据库名称
        list_discard, list_cmd_cp, list_import, result_cmd_xpt = data_uiaer(data_name_1)

        if not result_cmd_xpt:  # 如果为False则说明循环表的时候有错误
            judge_result_list = False

            file_error.write("------------" + data_name_1 + " 库------------" + "\n")
            file_error.write("mysql> use " + data_name_1 + ";" + "\n")
            file_error.write("mysql> SET Foreign_key_checks=0;" + "\n")
            file_error.write("mysql> " + "\n")
            for txt_file in list_discard:
                file_error.write(txt_file + "\n")

            file_error.write("shell> " + "\n")
            for cmd_cp in list_cmd_cp:
                file_error.write(cmd_cp + "\n")

            file_error.write("chown -R " + docker_userGroup + " " + mysql_datadir + "\n")
            file_error.write("mysql> " + "\n")
            for import_file in list_import:
                file_error.write(import_file + "\n")
            file_error.write("mysql> use " + data_name_1 + ";" + "\n")
            file_error.write("mysql> SET Foreign_key_checks=1;" + "\n")

    # 判断是否输出手动恢复的提示
    time_stamp = datetime.datetime.now().strftime('%Y.%m.%d-%H:%M:%S')
    if not judge_result_list:
        print("")
        print(time_stamp + "    \033[31m--->有导入错误的表,请手动执行 error_manualOperation.txt 中的命令<---\033[0m")

    print("")
    print(time_stamp + "    执行完毕,请重新连接数据库查看数据.")
